"""
/***************************************************************************
Bioidx
A QGIS plugin
Calculate the biodiversity index
                                                         -------------------
begin                                : 2009-08-03 
copyright                        : (C) 2009 by Mark Baas
email                                : mark.baas123@gmail.com 
 ***************************************************************************/

/***************************************************************************
 *                                                                                                                                                 *
 *     This program is free software; you can redistribute it and/or modify    *
 *     it under the terms of the GNU General Public License as published by    *
 *     the Free Software Foundation; either version 2 of the License, or         *
 *     (at your option) any later version.                                                                     *
 *                                                                                                                                                 *
 ***************************************************************************/
"""
# Import the PyQt and QGIS libraries
from PyQt4.QtCore import * 
from PyQt4.QtGui import *
from qgis.core import *
# Initialize Qt resources from file resources.py
import resources
# Import the code for the dialog
from BioidxDialog import BioidxDialog
import os, os.path

import math, tempfile
from distutils import file_util

def selectAll(pr):
    pr.select(pr.attributeIndexes(), QgsRectangle(), True)

class Bioidx: 

    def __init__(self, iface):
        # Save reference to the QGIS interface
        self.iface = iface

    def initGui(self):    
        # Create action that will start plugin configuration
        self.action = QAction(QIcon(":/plugins/bioidx/icon.png"), \
                "Simple biodiversity", self.iface.mainWindow())
        # connect the action to the run method
        QObject.connect(self.action, SIGNAL("triggered()"), self.run) 

        # Add toolbar button and menu item
        #self.iface.addToolBarIcon(self.action)
        #self.iface.addPluginToMenu("&Biodiversity index", self.action)
        menu_bar = self.iface.mainWindow().menuBar()
       
        self.ecogis_menu = None
        
        actions = menu_bar.actions()
        
        for action in actions:
            if action.text() == "EcoGis":
                self.ecogis_menu = action.menu()

        if not self.ecogis_menu: 
            self.ecogis_menu = QMenu()
            self.ecogis_menu.setTitle("EcoGis")
            menu_bar.insertMenu(actions[ len( actions ) - 1 ], self.ecogis_menu)

        self.bio_menu = None
            
        for action in self.ecogis_menu.actions():
            if action.text() == QCoreApplication.translate("EcoGis", "Biodiversity"):
                self.bio_menu = action.menu()
                
        if not self.bio_menu:
            self.bio_menu = QMenu()
            self.bio_menu.setTitle(QCoreApplication.translate("EcoGis", "Biodiversity"))
            self.ecogis_menu.addMenu(self.bio_menu)
        
        self.bio_menu.addActions([self.action])


    def unload(self):
        # Remove the plugin menu item and icon
        pass
        #self.iface.removePluginMenu("&Biodiversity index",self.action)
        #self.iface.removeToolBarIcon(self.action)

    # run method that performs all the real work
    def run(self): 
        # create and show the dialog 
        dlg = BioidxDialog(self.dlg_ok)
        # show the dialog
        dlg.show()
        result = dlg.exec_() 
        # See if OK was pressed

    def dlg_ok(self, dlg):
            if dlg == None:
                return
    
            # do something useful (delete the line containing pass and
            # substitute with your code
            
            vl_places = dlg.getLayer("places")
            vl_species = dlg.getLayer("species")
            hfield = dlg.getHField()
            efield = dlg.getEField()
            sourceField = dlg.getSourceField()
            distance = dlg.getDistance()
            #output = dlg.getOutputFilename()
            pb = dlg.ui.progressBar
            output = self.calculateDiversity(vl_places, vl_species, hfield, efield, sourceField, pb)

            # replace new layer for old layer
            

            
            #
            #addToTOC = QMessageBox.question(dlg, "Biodiversity", "Created output shapefile:" + "\n" + unicode(output) + "\n\nWould you like to add the new layer to the TOC? \n\nIf you choose to overwrite an already in the TOC shapefile, you have to remove the oldest one.", QMessageBox.Yes, QMessageBox.No, QMessageBox.NoButton)
            #if addToTOC == QMessageBox.Yes:
            self.addShapeToCanvas(unicode(output), vl_places)

    # Convinience function to add a vector layer to canvas based on input shapefile path ( as string )
    def addShapeToCanvas(self, shapeFilePath, layer):
        #shapeFilePathList = shapeFilePath.split( "/" )
        #layerName = QString( shapeFilePathList[len(shapeFilePathList)-1] )
        #if layerName.endsWith( ".shp" ):
        #    layerName = unicode( layerName ).rstrip( ".shp" )
        layerName = layer.name()
        vlayer_new = QgsVectorLayer( shapeFilePath, layerName, "ogr" )
        vlayer_new.copySymbologySettings(layer)

        if vlayer_new.isValid():
            QgsMapLayerRegistry.instance().addMapLayer(vlayer_new)
            QgsMapLayerRegistry.instance().removeMapLayer(layer.getLayerID())
            return True
        else:   
            return False

    def calculateDiversity(self, vl_places, vl_species, hfield, efield, sourceField, pb):
        # get providers
        pr_species = vl_species.dataProvider()
        pr_places = vl_places.dataProvider()
        
        # places = target layer
        output = pr_places.dataSourceUri().split("|")[0]
        output.replace("\\","/")

        # Select all
        selectAll(pr_species)
        selectAll(pr_places)

        # Create feature holders
        f_species = QgsFeature()
        f_places = QgsFeature()

        # add field to set H values
        fields = pr_places.fields()
        field = None

        for n in fields:
            o = fields[n]
            if o.name() == hfield:  field = n
        if not field:
            #print "field does not exists, adding it.."
            
            field = len(fields)
            fields[field] =  QgsField(hfield, QVariant.Double)
            fields[field+1] = QgsField(efield, QVariant.Double)
            
        # Setup a writer to write the new attributetabel
        if os.path.exists(output):
            # remove     
            QgsVectorFileWriter.deleteShapeFile(output)  

        writer = QgsVectorFileWriter(output, pr_places.encoding(),fields, pr_places.geometryType(), pr_places.crs() )

        # iterate through the features of places
        
        pb.setRange(0, pr_places.featureCount())
        pcnt = 0
        while pr_places.nextFeature(f_places):
            # init vars
            n_cases = 0
            species = {}
            
            # get the geometry
            geo_places = f_places.geometry()

            # Rewind to the first feature of species
            pr_species.rewind()

            while pr_species.nextFeature(f_species):
                # get the geometry
                geo_species = f_species.geometry()

                #if  geo_places.distance(geo_species) > distance: continue
                if not geo_places.intersects(geo_species): continue

                # get the attributes of this feature
                attrMap = f_species.attributeMap()    

                # now loop through the attributes
                for (i, attr) in attrMap.iteritems():
                    fieldName = pr_species.fields()[i].name()

                    if fieldName == sourceField:       
                        name = unicode(attr.toString())              
                        if not species.has_key(name):
                            species[name] = 0
                        
                        species[name] += 1
                        n_cases += 1

            # Calculate the H
            H = 0
            for specie in species:
                #print specie,"//", species[specie], "//"
                p = species[specie]/float(n_cases)
                H += p * math.log(p) # 

            H = H*-1
            if n_cases == 0: E = 0
            elif math.log(n_cases) == 0: E = 0
            else:
                E = H/math.log(n_cases)

            # get the attributes of this feature
            attrMap = f_places.attributeMap()  
            # store the H
            attrMap[field] = QVariant(H)
            attrMap[field+1] = QVariant(E)
            
            f_places.setAttributeMap(attrMap)

            writer.addFeature(f_places)
            
            pcnt += 1
            pb.setValue(pcnt)
        
        return output
        

